const Koa = require("koa");
const Router = require("@koa/router");
const sharp = require("sharp");
const tf = require("@tensorflow/tfjs");
const nsfwjs = require("nsfwjs");
const fs = require("fs");
const path = require("path");
const formidable = require("formidable");

const app = new Koa();
const router = new Router();

// 创建uploads目录（如果不存在）
if (!fs.existsSync(path.join(__dirname, "uploads"))) {
  fs.mkdirSync(path.join(__dirname, "uploads"));
}

// 声明模型变量，在服务器启动时加载
let nsfwModel = null;

// 加载NSFW模型
async function loadNSFWModel() {
  try {
    console.log("正在加载NSFW模型...");
    nsfwModel = await nsfwjs.load();
    console.log("NSFW模型加载成功！");
  } catch (error) {
    console.error("加载NSFW模型失败：", error);
    throw error;
  }
}

// 图像分类函数
async function classifyImage(imageBuffer) {
  try {
    // 使用sharp处理图像，确保尺寸正确
    const processedImage = await sharp(imageBuffer)
      .resize(224, 224, { fit: "inside" })
      .toBuffer();

    // 使用sharp获取原始图像数据
    const { data, info } = await sharp(processedImage).raw().toBuffer({
      resolveWithObject: true,
    });

    // 创建张量
    const tensor = tf.tensor3d(new Uint8Array(data), [
      info.height,
      info.width,
      info.channels,
    ]);

    // 进行分类
    const predictions = await nsfwModel.classify(tensor);

    // 清理张量
    tensor.dispose();

    return predictions;
  } catch (error) {
    console.error("图像分类失败：", error);
    throw error;
  }
}

// 定义/checkImg接口
router.post("/checkImg", async (ctx) => {
  // 使用formidable处理文件上传
  const form = new formidable.IncomingForm({
    uploadDir: path.join(__dirname, "uploads"),
    keepExtensions: true,
    maxFileSize: 5 * 1024 * 1024, // 5MB
    multiples: false,
  });

  // 将formidable转换为Promise
  return new Promise((resolve, reject) => {
    form.parse(ctx.req, async (err, fields, files) => {
      try {
        if (err) {
          ctx.status = 400;
          ctx.body = {
            status: false,
            reason: "文件上传失败：" + err.message,
          };
          return resolve();
        }

        // 检查是否有文件上传
        if (!files || !files.image) {
          // 添加调试信息
          console.log("Files structure:", files);
          ctx.status = 400;
          ctx.body = {
            status: false,
            reason: "请上传图像文件",
          };
          return resolve();
        }

        const image = files.image;

        // 打印文件对象的详细结构，用于调试
        console.log("Image file structure:", JSON.stringify(image, null, 2));

        // 获取正确的文件路径 - 尝试多种可能的属性名
        const filePath = image[0].filepath;
        console.log("Image file structure2:", filePath);
        if (!filePath) {
          ctx.status = 400;
          ctx.body = {
            status: false,
            reason: "无法获取文件路径",
          };
          return resolve();
        }

        console.log("Using file path:", filePath);

        try {
          const imageBuffer = fs.readFileSync(filePath);

          // 进行图像分类
          const predictions = await classifyImage(imageBuffer);
          console.log("Predictions:", predictions);
          // 检查是否包含不良内容
          const adultContent = predictions.find(
            (pred) =>
              (pred.className === "Porn" && pred.probability > 0.6) ||
              (pred.className === "Hentai" && pred.probability > 0.5) ||
              (pred.className === "Sexy" && pred.probability > 0.7)
          );

          if (adultContent) {
            ctx.body = {
              status: false,
              reason: `检测到${adultContent?.className}内容，概率：${(
                adultContent?.probability * 100
              ).toFixed(2)}%`,
            };
          } else {
            ctx.body = {
              status: true,
            };
          }

          // 清理临时文件
          fs.unlinkSync(filePath);
        } catch (fileError) {
          console.error("文件处理失败：", fileError);
          ctx.status = 500;
          ctx.body = {
            status: false,
            reason: "文件处理失败：" + fileError.message,
          };
        }

        resolve();
      } catch (error) {
        console.error("处理请求失败：", error);
        ctx.status = 500;
        ctx.body = {
          status: false,
          reason: "服务器内部错误：" + error.message,
        };
        resolve();
      }
    });
  });
});

// 挂载路由
app.use(router.routes());
app.use(router.allowedMethods());

// 启动服务器
async function startServer() {
  try {
    // 先加载模型
    await loadNSFWModel();

    // 启动Koa服务器
    const PORT = process.env.PORT || 3001;
    app.listen(PORT, () => {
      console.log(`服务器运行在 http://localhost:${PORT}`);
      console.log("API接口：POST /checkImg");
    });
  } catch (error) {
    console.error("启动服务器失败：", error);
  }
}

// 启动服务器
startServer();
